"""
Agentic RAG (Retrieval-Augmented Generation) 智能代理测试页面

本模块实现了基于 LangGraph 的智能RAG系统，具备以下功能：
- 自适应检索决策：根据问题智能决定是否需要检索文档
- 文档质量评估：对检索到的文档进行相关性评分
- 问题重写优化：对无关检索结果时自动重写问题
- 智能回答生成：基于上下文生成精准回答

注意：当前版本包含硬编码的API密钥，生产环境中应使用配置文件管理
"""
import streamlit as st
from langchain.tools.retriever import create_retriever_tool
from langgraph.graph import MessagesState
from langchain.chat_models import init_chat_model
import os
from pydantic import BaseModel, Field
from typing import Literal, List
from langgraph.graph import StateGraph, START, END
from langgraph.prebuilt import ToolNode
from langgraph.prebuilt import tools_condition
# 导入项目服务
from service.vector_service import (
    get_vector_service, VectorService,
    get_all_collections, get_collection_names, 
    get_collection_details, get_collections_summary
)
from service.logger_service import get_logger_service


# ======================== 全局变量 ========================
logger = get_logger_service()

# ======================== 集合管理功能 ========================
# 这些功能现在直接使用 vector_service.py 中的函数

def get_available_collections() -> List[dict]:
    """
    获取可用的向量集合列表（使用 vector_service）
    
    Returns:
        List[dict]: 包含集合详细信息的列表
    """
    return get_all_collections()

def get_collection_info(collection_name: str) -> dict:
    """
    获取指定集合的详细信息（使用 vector_service）
    
    Args:
        collection_name: 集合名称
        
    Returns:
        dict: 集合详细信息
    """
    return get_collection_details(collection_name)

def get_selected_vector_service(collection_name: str) -> VectorService:
    """
    根据选择的集合获取向量服务
    Args:
        collection_name: 集合名称
    Returns:
        VectorService: 向量服务实例
    """
    return get_vector_service(
        collection_name=collection_name,
        embeddings_model="text-embedding-v3",
        distance_strategy="cosine"
    )

# ======================== 语言模型配置 ========================
@st.cache_resource
def initialize_language_models():
    """
    初始化语言模型，使用缓存避免重复初始化
    """
    try:
        # 设置DeepSeek API密钥
        os.environ["DEEPSEEK_API_KEY"] = st.secrets["deepseek"]["key"]
        # 初始化响应模型和评估模型
        response_model = init_chat_model("deepseek-chat", model_provider="deepseek")
        grader_model = init_chat_model("deepseek-chat", model_provider="deepseek")
        return response_model, grader_model
    except Exception as e:
        st.error(f"❌ 语言模型初始化失败：{str(e)}")
        return None, None

# 初始化语言模型
response_model, grader_model = initialize_language_models()
if not response_model or not grader_model:
    st.stop()

# ======================== 动态检索工具创建 ========================
def create_dynamic_retriever_tool(vector_service: VectorService, collection_name: str):
    """
    根据向量服务动态创建检索工具
    Args:
        vector_service: 向量服务实例
        collection_name: 集合名称
    Returns:
        检索工具
    """
    # 创建检索器，使用MMR算法提升多样性
    retriever = vector_service.vector_store.as_retriever(
        search_type="mmr", 
        search_kwargs={"k": 3}  # 增加检索数量以提高答案质量
    )
    # 创建检索工具
    tool = create_retriever_tool(
        retriever,
        "retrieve_documents",
        f"搜索并返回来自 '{collection_name}' 集合的相关文档信息，用于回答用户问题。",
    )
    return tool

# ======================== 智能代理节点函数 ========================
def create_agent_functions(retriever_tool):
    """
    创建智能代理节点函数，避免全局变量依赖
    Args:
        retriever_tool: 检索工具实例
    Returns:
        tuple: (generate_query_or_respond函数, grade_documents函数, rewrite_question函数, generate_answer函数)
    """
    def generate_query_or_respond(state: MessagesState):
        """
        智能查询生成或直接响应节点
        根据当前对话状态，模型将智能决定：
        1. 使用检索工具搜索相关文档
        2. 直接基于已有知识回答用户问题
        Args:
            state (MessagesState): 当前对话状态，包含历史消息
        Returns:
            dict: 包含模型响应的消息状态更新
        """
        response = (
            response_model
            .bind_tools([retriever_tool]).invoke(state["messages"])
        )
        return {"messages": [response]}
    
    def grade_documents(state: MessagesState) -> Literal["generate_answer", "rewrite_question"]:
        """
        文档相关性评估节点
        评估检索到的文档是否与用户问题相关：
        - 相关：进入答案生成流程
        - 不相关：进入问题重写流程
        Args:
            state (MessagesState): 当前对话状态
        Returns:
            Literal: 下一个节点的名称
        """
        question = state["messages"][0].content
        context = state["messages"][-1].content
        prompt = GRADE_PROMPT.format(question=question, context=context)
        response = (
            grader_model
            .with_structured_output(GradeDocuments).invoke(
                [{"role": "user", "content": prompt}]
            )
        )
        score = response.binary_score
        if score == "yes":
            return "generate_answer"
        else:
            return "rewrite_question"
    
    def rewrite_question(state: MessagesState):
        """
        问题重写优化节点
        当检索到的文档与问题不相关时，重新表述问题以提高检索精度。
        通过理解问题的语义意图，生成更精准的查询表达。
        Args:
            state (MessagesState): 当前对话状态
        Returns:
            dict: 包含重写后问题的消息状态更新
        """
        messages = state["messages"]
        question = messages[0].content
        prompt = REWRITE_PROMPT.format(question=question)
        response = response_model.invoke([{"role": "user", "content": prompt}])
        return {"messages": [{"role": "user", "content": response.content}]}
    
    def generate_answer(state: MessagesState):
        """
        智能回答生成节点
        基于检索到的相关文档上下文，生成准确、简洁的答案。
        遵循诚实性原则：信息不足时承认不知道。
        Args:
            state (MessagesState): 当前对话状态
        Returns:
            dict: 包含生成答案的消息状态更新
        """
        question = state["messages"][0].content
        context = state["messages"][-1].content
        prompt = GENERATE_PROMPT.format(question=question, context=context)
        response = response_model.invoke([{"role": "user", "content": prompt}])
        return {"messages": [response]}
    
    return generate_query_or_respond, grade_documents, rewrite_question, generate_answer


# ======================== 提示词模板定义 ========================

# 文档相关性评估提示词模板（中文优化版）
GRADE_PROMPT = (
    "你是一个文档相关性评估专家，需要判断检索到的文档是否与用户问题相关。\n"
    "检索到的文档内容：\n\n{context}\n\n"
    "用户问题：{question}\n\n"
    "如果文档包含与用户问题相关的关键词或语义内容，则判定为相关。\n"
    "请给出二元评分：'yes'表示相关，'no'表示不相关。"
)

# 问题重写提示词模板（中文优化版）
REWRITE_PROMPT = (
    "请分析输入问题的潜在语义意图和含义。\n"
    "原始问题：\n"
    "------- \n"
    "{question}"
    "\n ------- \n"
    "请重新表述一个更清晰、更具体的改进问题："
)

# 回答生成提示词模板（中文优化版）
GENERATE_PROMPT = (
    "你是一个专业的问答助手。请基于提供的检索上下文回答问题。\n"
    "如果上下文中没有足够信息回答问题，请诚实地说不知道。\n"
    "请保持回答简洁明了，最多使用三句话。\n\n"
    "问题：{question}\n"
    "上下文：{context}"
)

class GradeDocuments(BaseModel):
    """文档相关性评估模型"""
    binary_score: str = Field(
        description="相关性评分：'yes'表示相关，'no'表示不相关"
    )

# ======================== 智能工作流构建 ========================

def create_agentic_workflow(retriever_tool):
    """
    创建智能RAG工作流
    
    Args:
        retriever_tool: 检索工具
        
    Returns:
        编译后的工作流图
    """
    # 创建代理函数
    generate_query_or_respond, grade_documents, rewrite_question, generate_answer = create_agent_functions(retriever_tool)
    
    # 初始化状态图工作流
    workflow = StateGraph(MessagesState)

    # 添加核心节点到工作流
    workflow.add_node("generate_query_or_respond", generate_query_or_respond)  # 智能查询/响应节点
    workflow.add_node("retrieve", ToolNode([retriever_tool]))  # 文档检索节点
    workflow.add_node("rewrite_question", rewrite_question)  # 问题重写节点
    workflow.add_node("generate_answer", generate_answer)  # 答案生成节点

    # 设置工作流入口点
    workflow.add_edge(START, "generate_query_or_respond")

    # ======================== 条件边配置 ========================

    # 智能路由：决定是否进行文档检索
    workflow.add_conditional_edges(
        "generate_query_or_respond",
        # 评估大模型决策（调用检索工具或直接响应用户）
        tools_condition,
        {
            # 条件输出到图节点的映射
            "tools": "retrieve",  # 需要工具 -> 进入检索流程
            END: END,            # 无需工具 -> 直接结束
        },
    )

    # 文档相关性评估路由
    workflow.add_conditional_edges(
        "retrieve",
        # 评估检索文档的相关性
        grade_documents,
        # 返回值直接映射到节点名称：
        # "generate_answer" -> 文档相关，生成答案
        # "rewrite_question" -> 文档不相关，重写问题
    )

    # ======================== 静态边配置 ========================
    workflow.add_edge("generate_answer", END)  # 生成答案后结束
    workflow.add_edge("rewrite_question", "generate_query_or_respond")  # 重写问题后重新开始

    # 编译工作流图
    return workflow.compile()

# ======================== Streamlit 测试界面 ========================

def main():
    """
    Streamlit 主界面
    提供Agentic RAG系统的交互式测试环境
    """
    st.set_page_config(
        page_title="Agentic RAG 智能代理测试",
        page_icon="🤖",
        layout="wide"
    )
    
    st.title("🤖 Agentic RAG 智能代理测试系统")
    st.markdown("""
    这是一个基于 LangGraph 的智能RAG系统测试界面。系统具备以下能力：
    
    - **智能决策**：自动判断是否需要检索文档
    - **质量评估**：对检索结果进行相关性评分  
    - **自我优化**：不相关时自动重写查询问题
    - **精准回答**：基于上下文生成准确答案
    - **多集合支持**：可选择不同知识库进行检索
    """)
    
    # 侧边栏配置
    with st.sidebar:
        st.header("🔧 系统配置")
        
        # 集合选择
        st.subheader("📚 知识库选择")
        
        # 获取集合信息
        collections_data = get_available_collections()
        collection_names = [col["name"] for col in collections_data]
        
        if not collection_names:
            st.error("❌ 未找到任何集合")
            return
            
        # 显示集合选择器，并显示文档数量
        collection_options = []
        for col in collections_data:
            status_icon = "📁" if col["doc_count"] > 0 else "📂"
            option_text = f"{status_icon} {col['name']} ({col['doc_count']} 文档)"
            collection_options.append(option_text)
        
        selected_index = st.selectbox(
            "选择知识库集合：",
            range(len(collection_options)),
            format_func=lambda x: collection_options[x],
            index=0,
            help="选择要用于检索的向量集合"
        )
        
        selected_collection = collection_names[selected_index]
        collection_data = collections_data[selected_index]
        
        # 获取选中集合的向量服务
        if selected_collection:
            vector_service = get_selected_vector_service(selected_collection)
            
            # 显示当前配置
            st.subheader("当前配置")
            
            # 显示集合状态
            if collection_data["is_empty"]:
                st.warning(f"⚠️ 集合 '{selected_collection}' 为空")
            else:
                st.success(f"✅ 集合 '{selected_collection}' 包含 {collection_data['doc_count']} 个文档")
            
            # 配置信息
            st.info(f"""
            **知识库**: {selected_collection}
            **集合ID**: {collection_data.get('uuid', 'N/A')[:8]}...
            **文档数量**: {collection_data['doc_count']}
            **嵌入模型**: {vector_service.embeddings_model}
            **语言模型**: DeepSeek Chat
            **检索策略**: MMR (k=3)
            **距离策略**: {vector_service.distance_strategy.value}
            """)
            
            # 显示集合元数据（如果有）
            if collection_data.get("metadata"):
                with st.expander("📋 集合元数据"):
                    st.json(collection_data["metadata"])
        
        # 系统状态检查
        st.subheader("📊 系统状态")
        if st.button("检查连接状态"):
            if selected_collection:
                try:
                    # 测试向量存储连接
                    test_results = vector_service.similarity_search("test", k=1)
                    st.success("✅ 向量存储连接正常")
                    
                    # 测试模型连接
                    model_test = response_model.invoke("测试连接")
                    st.success("✅ 语言模型连接正常")
                    
                    # 显示服务状态
                    status = vector_service.get_service_status()
                    st.json(status)
                    
                except Exception as e:
                    st.error(f"❌ 连接测试失败: {str(e)}")
            else:
                st.warning("⚠️ 请先选择知识库集合")
    
    # 检查是否选择了集合
    if not selected_collection:
        st.warning("⚠️ 请在左侧边栏选择一个知识库集合")
        return
        
    # 创建动态检索工具和工作流
    retriever_tool = create_dynamic_retriever_tool(vector_service, selected_collection)
    workflow_graph = create_agentic_workflow(retriever_tool)
    
    # 主界面分为两列
    col1, col2 = st.columns([1, 1])
    
    with col1:
        st.subheader("💬 智能问答测试")
        
        # 问题输入
        user_question = st.text_area(
            "请输入您的问题：",
            placeholder="例如：什么是强化学习中的奖励黑客攻击？",
            height=100
        )
        
        # 测试按钮
        if st.button("🚀 开始智能问答", type="primary"):
            if user_question.strip():
                with st.spinner("🤖 智能代理正在处理您的问题..."):
                    try:
                        # 调用智能工作流
                        result = workflow_graph.invoke({"messages": [{"role": "user", "content": user_question}]})
                        
                        # 显示结果
                        if result and "messages" in result:
                            final_response = result["messages"][-1]
                            
                            st.success("✅ 处理完成！")
                            st.markdown("### 🎯 智能回答")
                            
                            # 根据响应类型显示不同内容
                            if hasattr(final_response, 'content'):
                                st.markdown(final_response.content)
                            else:
                                st.markdown(str(final_response))
                                
                            # 记录用户活动
                            logger.log_system_event(
                                "agentic_rag_query", 
                                f"用户查询: '{user_question}', 集合: {selected_collection}"
                            )
                        else:
                            st.warning("⚠️ 未收到有效响应")
                            
                    except Exception as e:
                        st.error(f"❌ 处理过程中出现错误：{str(e)}")
                        logger.log_system_event("agentic_rag_error", f"查询失败: {e}", level="error")
            else:
                st.warning("⚠️ 请输入问题后再提交")
    
    with col2:
        st.subheader("📈 处理流程追踪")
        
        if st.button("🔍 显示工作流图结构"):
            st.markdown("""
            ### 🔄 智能工作流程图
            
            ```
            开始 → 智能查询/响应
                    ↓
            ┌─ 需要检索? ─┐
            │              │
            ↓ 是          │ 否 → 结束
            文档检索        
            ↓              
            文档相关性评估   
            ↓              
            ┌─ 相关? ─┐     
            │         │     
            ↓ 是      ↓ 否  
            生成答案   重写问题
            ↓         ↓     
            结束      回到开始
            ```
            """)
        
        st.subheader("🎯 快速测试用例")
        
        test_cases = [
            "什么是机器学习？",
            "深度学习和传统机器学习的区别是什么？", 
            "如何评估模型的性能？",
            "什么是过拟合现象？",
            "RAG系统的工作原理是什么？",
            "向量数据库有什么优势？"
        ]
        
        st.markdown("**点击以下问题快速测试：**")
        for i, case in enumerate(test_cases):
            if st.button(f"📝 {case}", key=f"test_case_{i}"):
                st.code(case)
                st.info("点击上方 '开始智能问答' 按钮测试此问题")
                
        # 集合管理工具
        st.subheader("🛠️ 集合管理工具")
        
        # 显示所有集合的统计信息
        if st.button("📊 查看所有集合统计"):
            st.markdown("### 📈 集合统计信息")
            
            # 使用 vector_service 的汇总功能
            summary = get_collections_summary()
            
            # 显示汇总指标
            col1, col2, col3, col4 = st.columns(4)
            with col1:
                st.metric("总集合数", summary["total_collections"])
            with col2:
                st.metric("总文档数", summary["total_documents"])
            with col3:
                st.metric("非空集合", summary["non_empty_collections"])
            with col4:
                st.metric("空集合", summary["empty_collections"])
            
            # 创建详细统计表格
            stats_data = []
            for col in collections_data:
                stats_data.append({
                    "集合名称": col["name"],
                    "文档数量": col["doc_count"],
                    "状态": "有文档" if col["doc_count"] > 0 else "空集合",
                    "集合ID": col.get("uuid", "N/A")[:8] + "..." if col.get("uuid") else "N/A"
                })
            
            import pandas as pd
            stats_df = pd.DataFrame(stats_data)
            st.dataframe(stats_df, use_container_width=True)
            
            # 显示集合名称列表
            if summary["collection_names"]:
                st.markdown("**所有集合名称:**")
                st.write(", ".join(summary["collection_names"]))
        
        if st.button("📋 查看当前集合详情"):
            if selected_collection and not collection_data["is_empty"]:
                with st.expander(f"集合 '{selected_collection}' 详情", expanded=True):
                    # 获取文档样例
                    docs_df = vector_service.get_documents(limit=3)
                    if not docs_df.empty:
                        st.markdown("**文档样例:**")
                        for idx, row in docs_df.iterrows():
                            with st.container():
                                st.markdown(f"**文档 {row['id']}:**")
                                content = row['document']
                                if len(content) > 300:
                                    content = content[:300] + "..."
                                st.text(content)
                                
                                # 显示元数据
                                if row.get('cmetadata'):
                                    st.json(row['cmetadata'])
                                st.divider()
                    else:
                        st.info("该集合暂无文档")
                        
                    # 获取元数据统计
                    metadata_stats = vector_service.get_metadata_stats()
                    if not metadata_stats.empty:
                        st.markdown("**元数据统计:**")
                        st.dataframe(metadata_stats, use_container_width=True)
            elif collection_data["is_empty"]:
                st.warning(f"集合 '{selected_collection}' 为空，无法显示详情")
            else:
                st.warning("请先选择一个集合")
        
        # 相似度搜索测试
        st.markdown("---")
        st.markdown("**🔍 相似度搜索测试**")
        
        test_query = st.text_input(
            "输入测试查询：",
            value="机器学习",
            key="search_test_query",
            help="输入查询文本来测试向量搜索"
        )
        
        search_k = st.slider("返回结果数量：", 1, 10, 3, key="search_k")
        
        if st.button("🚀 执行搜索", type="primary"):
            if selected_collection and test_query.strip():
                if not collection_data["is_empty"]:
                    try:
                        with st.spinner("🔍 搜索中..."):
                            search_results = vector_service.similarity_search_with_score(test_query, k=search_k)
                        
                        if search_results:
                            st.markdown("### 🎯 搜索结果")
                            for i, result in enumerate(search_results):
                                with st.expander(f"结果 {i+1} - 相似度: {result.score:.4f}", expanded=i==0):
                                    st.markdown("**内容:**")
                                    st.text(result.content)
                                    
                                    if result.metadata:
                                        st.markdown("**元数据:**")
                                        st.json(result.metadata)
                        else:
                            st.info("未找到相关结果")
                            
                    except Exception as e:
                        st.error(f"搜索失败: {str(e)}")
                else:
                    st.warning(f"集合 '{selected_collection}' 为空，无法执行搜索")
            elif not test_query.strip():
                st.warning("请输入查询文本")
            else:
                st.warning("请先选择一个集合")

if __name__ == "__main__":
    main()

# ======================== 文档预处理示例（注释保留） ========================
"""
以下代码展示了如何预处理文档并添加到向量存储中：

from langchain_community.document_loaders import WebBaseLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter

# 示例：加载 Lilian Weng 博客文章
urls = [
    "https://lilianweng.github.io/posts/2024-11-28-reward-hacking/",
    "https://lilianweng.github.io/posts/2024-07-07-hallucination/",
    "https://lilianweng.github.io/posts/2024-04-12-diffusion-video/",
]

# 加载文档
docs = [WebBaseLoader(url).load() for url in urls]

# 文档分割
docs_list = [item for sublist in docs for item in sublist]
text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
    chunk_size=100, chunk_overlap=50
)
doc_splits = text_splitter.split_documents(docs_list)

# 添加到向量存储
vector_store.add_documents(doc_splits)
"""